{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import scipy, sys, os, pyproj, glob, re, h5py\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from scipy.signal import correlate\n",
    "from astropy.time import Time\n",
    "\n",
    "# Import some of the scripts that we have written\n",
    "import sys\n",
    "sys.path.append(\"/home/jovyan/surface_velocity/scripts\")\n",
    "from loading_scripts import atl06_to_dict\n",
    "\n",
    "%matplotlib widget\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190430122344_04920311_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20181030210407_04920111_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190730080323_04920411_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190220230230_08320211_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190312235510_11380211_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20181108184743_06280111_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190228224553_09540211_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190623094402_13150311_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190506112405_05830311_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190611193446_11380311_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190620171809_12740311_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190823150456_08630411_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190529105941_09340311_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190308130329_10700211_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190824143917_08780411_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190812071246_06900411_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20181118033101_07710111_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190822060451_08420411_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190607194306_10770311_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "file /home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190129164404_04920211_003_01.h5 encountered error 'Unable to open object (component not found)'\n",
      "read 613 data files of which 20 gave errors\n"
     ]
    }
   ],
   "source": [
    "datapath = '/home/jovyan/shared/surface_velocity/FIS_ATL06/'\n",
    "\n",
    "# find all the files in the directory:\n",
    "# ATL06_files=glob.glob(os.path.join(datapath, 'PIG_ATL06', '*.h5'))\n",
    "\n",
    "ATL06_files=glob.glob(os.path.join(datapath, '*.h5'))\n",
    "\n",
    "D_dict={}\n",
    "error_count=0\n",
    "for file in ATL06_files:\n",
    "    try:\n",
    "        D_dict[file]=atl06_to_dict(file, '/gt2l', index=slice(0, -1, 25), epsg=3031)\n",
    "    except KeyError as e:\n",
    "        print(f'file {file} encountered error {e}')\n",
    "        error_count += 1\n",
    "print(f\"read {len(D_dict)} data files of which {error_count} gave errors\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4 2\n",
      "filename=/home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190509193251_06340311_003_01.h5\n",
      "filename=/home/jovyan/shared/surface_velocity/FIS_ATL06/processed_ATL06_20190808151232_06340411_003_01.h5\n"
     ]
    },
    {
     "ename": "KeyError",
     "evalue": "'gt1l'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-11-ffc3d9211227>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     99\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    100\u001b[0m     \u001b[0mx1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m2.915e7\u001b[0m\u001b[0;31m#x_atc[cycles[0]][beams[0]][1000] <-- the very first x value in a file; doesn't work, I think b/c nans # 2.93e7\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 101\u001b[0;31m     \u001b[0mx1s\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx_atc\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mcycles\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mveloc_number\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mbeams\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0msearch_width\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;36m500\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0msegment_length\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0msearch_width\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    102\u001b[0m     \u001b[0mvelocities\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    103\u001b[0m     \u001b[0mx_out\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mempty_like\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx1s\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyError\u001b[0m: 'gt1l'"
     ]
    }
   ],
   "source": [
    "beams = ['gt1l']#,'gt1r','gt2l','gt2r','gt3l','gt3r']\n",
    "\n",
    "for i,key in enumerate(D_dict.keys()):\n",
    "    \n",
    "    rgt=key[78:82]\n",
    "\n",
    "    D_2l={}\n",
    "    D_2r={}\n",
    "\n",
    "    # iterate over the repeat cycles\n",
    "    for cycle in ['03','04','05','06','07']:\n",
    "        for filename in glob.glob(os.path.join(datapath, f'*ATL06_*_{rgt}{cycle}*_003*.h5')):\n",
    "            try:\n",
    "                # read the left-beam data\n",
    "                D_2l[filename]=atl06_to_dict(filename,'/gt2l', index=None, epsg=3031)\n",
    "                # read the right-beam data\n",
    "                D_2r[filename]=atl06_to_dict(filename,'/gt2r', index=None, epsg=3031)\n",
    "            except Exception as e:\n",
    "                print(f'filename={filename}, exception={e}')\n",
    "\n",
    "\n",
    "    cycles = [] # names of cycles with data\n",
    "    for filename, Di in D_2l.items():\n",
    "        cycles += [str(Di['cycle']).zfill(2)]\n",
    "    cycles.sort()\n",
    "\n",
    "    if len(cycles)<2:\n",
    "        continue\n",
    "    print(i,len(cycles))\n",
    "    \n",
    "    ### extract and plot data from all available cycles\n",
    "    x_atc = {}\n",
    "    h_li = {}\n",
    "    h_li_diff = {}\n",
    "    times = {}\n",
    "    for cycle in cycles:\n",
    "        # find Di that matches cycle:\n",
    "        Di = {}\n",
    "        x_atc[cycle] = {}\n",
    "        h_li[cycle] = {}\n",
    "        h_li_diff[cycle] = {}\n",
    "        times[cycle] = {}\n",
    "\n",
    "        filenames = glob.glob(os.path.join(datapath, f'*ATL06_*_{rgt}{cycle}*_003*.h5'))\n",
    "        for filename in filenames:\n",
    "            try:\n",
    "                for beam in beams:\n",
    "                    Di[filename]=atl06_to_dict(filename,'/'+ beam, index=None, epsg=3031)\n",
    "\n",
    "                    times[cycle][beam] = Di[filename]['data_start_utc']\n",
    "\n",
    "                    # extract h_li and x_atc for that section\n",
    "                    x_atc_tmp = Di[filename]['x_atc']\n",
    "                    h_li_tmp = Di[filename]['h_li']#[ixs]\n",
    "\n",
    "                    # segment ids:\n",
    "                    seg_ids = Di[filename]['segment_id']\n",
    "    #                 print(len(seg_ids), len(x_atc_tmp))\n",
    "\n",
    "                    # make a monotonically increasing x vector\n",
    "                    # assumes dx = 20 exactly, so be carefull referencing back\n",
    "                    ind = seg_ids - np.nanmin(seg_ids) # indices starting at zero, using the segment_id field, so any skipped segment will be kept in correct location\n",
    "                    x_full = np.arange(np.max(ind)+1) * 20 + x_atc_tmp[0]\n",
    "                    h_full = np.zeros(np.max(ind)+1) + np.NaN\n",
    "                    h_full[ind] = h_li_tmp\n",
    "\n",
    "\n",
    "                    x_atc[cycle][beam] = x_full\n",
    "                    h_li[cycle][beam] = h_full\n",
    "\n",
    "\n",
    "    #                 ### here is where you would put a filter\n",
    "    #                 # you would want to de-mean and detrend that section first:\n",
    "    #                 h = h_full\n",
    "    #                 x = x_full\n",
    "    #                 h = h - np.nanmean(h) # de-mean\n",
    "    #                 h = scipy.signal.detrend(h, type = 'linear') # de-trend; need to deal with nans first\n",
    "    #                 # use scipy.signal.filter to filter\n",
    "\n",
    "    #                 # differentiate that section of data\n",
    "                    h_diff = (h_full[1:] - h_full[0:-1]) / (x_full[1:] - x_full[0:-1])\n",
    "                    h_li_diff[cycle][beam] = h_diff\n",
    "\n",
    "            except:\n",
    "                print(f'filename={filename}')\n",
    "                \n",
    "                \n",
    "    n_veloc = len(cycles) - 1\n",
    "    veloc_number = 0\n",
    "\n",
    "    segment_length = 3000 # m\n",
    "\n",
    "    search_width = 500 # m\n",
    "    dx = 20 # meters between x_atc points\n",
    "\n",
    "    x1 = 2.915e7#x_atc[cycles[0]][beams[0]][1000] <-- the very first x value in a file; doesn't work, I think b/c nans # 2.93e7\n",
    "    x1s = x_atc[cycles[veloc_number]][beams[0]][search_width+500:-segment_length-2*search_width:10]\n",
    "    velocities = {}  \n",
    "    x_out = np.empty_like(x1s)\n",
    "    y_out = np.empty_like(x1s)\n",
    "    for beam in beams:\n",
    "        velocities[beam] = np.empty_like(x1s)\n",
    "    for xi,x1 in enumerate(x1s):\n",
    "        for veloc_number in range(n_veloc):\n",
    "            cycle1 = cycles[veloc_number]\n",
    "            cycle2 = cycles[veloc_number+1]\n",
    "            t1_string = times[cycle1]['gt1l'][0].astype(str) #figure out later if just picking hte first one it ok\n",
    "            t1 = Time(t1_string)\n",
    "\n",
    "            t2_string = times[cycle2]['gt1l'][0].astype(str) #figure out later if just picking hte first one it ok\n",
    "            t2 = Time(t2_string)\n",
    "\n",
    "            dt = (t2 - t1).jd # difference in julian days\n",
    "\n",
    "            for beam in beams:\n",
    "                # cut out small chunk of data at time t1 (first cycle)\n",
    "                x_full_t1 = x_atc[cycle1][beam]\n",
    "                ix_x1 = np.arange(len(x_full_t1))[x_full_t1 >= x1][0]\n",
    "                ix_x2 = ix_x1 + int(np.round(segment_length/dx))      \n",
    "                x_t1 = x_full_t1[ix_x1:ix_x2]\n",
    "                h_li1 = h_li_diff[cycle1][beam][ix_x1-1:ix_x2-1] # start 1 index earlier because the data are differentiated\n",
    "\n",
    "                # cut out a wider chunk of data at time t2 (second cycle)\n",
    "                x_full_t2 = x_atc[cycle2][beam]\n",
    "                ix_x3 = ix_x1 - int(np.round(search_width/dx)) # offset on earlier end by # indices in search_width\n",
    "                ix_x4 = ix_x2 + int(np.round(search_width/dx)) # offset on later end by # indices in search_width\n",
    "                x_t2 = x_full_t2[ix_x3:ix_x4]\n",
    "                h_li2 = h_li_diff[cycle2][beam][ix_x3:ix_x4]\n",
    "                \n",
    "                try:\n",
    "                    # correlate old with newer data\n",
    "                    corr = correlate(h_li1, h_li2, mode = 'valid', method = 'direct') \n",
    "                    norm_val = np.sqrt(np.sum(h_li1**2)*np.sum(h_li2**2)) # normalize so values range between 0 and 1\n",
    "                    corr = corr / norm_val\n",
    "                except:\n",
    "                    continue\n",
    "\n",
    "\n",
    "        #         lagvec = np.arange( -(len(h_li1) - 1), len(h_li2), 1)# for mode = 'full'\n",
    "        #         lagvec = np.arange( -int(search_width/dx) - 1, int(search_width/dx) +1, 1) # for mode = 'valid'\n",
    "                lagvec = np.arange(- int(np.round(search_width/dx)), int(search_width/dx) +1,1)# for mode = 'valid'\n",
    "\n",
    "                shift_vec = lagvec * dx\n",
    "\n",
    "                if all(np.isnan(corr)):\n",
    "                    velocities[beam][xi] = np.nan\n",
    "                else:\n",
    "                    ix_peak = np.arange(len(corr))[corr == np.nanmax(corr)][0]\n",
    "                    best_lag = lagvec[ix_peak]\n",
    "                    best_shift = shift_vec[ix_peak]\n",
    "                    velocities[beam][xi] = best_shift/(dt/365)\n",
    "                \n",
    "                try:\n",
    "                    x_out[xi] = Di[filename]['x'][(ix_x1+ix_x2)//2]\n",
    "                    y_out[xi] = Di[filename]['y'][(ix_x1+ix_x2)//2]\n",
    "                except:\n",
    "                    x_out[xi] = np.nan\n",
    "                    y_out[xi] = np.nan\n",
    "                    \n",
    "    out = np.array([x_out,y_out,velocities[beam]])\n",
    "    np.save('/home/jovyan/shared/surface_velocity/Corr_Vel/corr_vel_'+rgt,out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'03': {}, '04': {}}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_atc"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
